import asyncio
import json
import os.path
from asyncio.locks import Lock
import asyncpg
import datetime
from typing import *
from loguru import logger
from collections import defaultdict
from config import DATABASE


class Checker:
    host = DATABASE['postgresql']['host']
    port = DATABASE['postgresql']['port']
    user = DATABASE['postgresql']['username']
    password = DATABASE['postgresql']['password']
    db = DATABASE['postgresql']['db']

    tables = {
        60: 'kline_minute',
        3600: 'kline_hour',
        86400: 'kline_day',
    }

    def __init__(self, start: int, end: int, path: str = None, symbols: List[str] = None):
        self.symbols: Optional[List[str]] = symbols
        self.start = start
        self.end = end
        self.pool = None
        self.path = path
        self.lock = Lock()
        self.info: Dict[int, Dict[str, dict]] = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))

    async def get_conn(self) -> asyncpg.Connection:
        if self.pool is None:
            self.pool = await asyncpg.create_pool(host=self.host, port=self.port, user=self.user,
                                                  password=self.password,
                                                  database=self.db, min_size=20, max_size=100)
        return await self.pool.acquire()

    async def check_symbol(self, period: int, symbol: str, trade_type: int, daily: bool = True, onboard: int = 0):
        """
        :param onboard: 币种上线时间
        :param daily:
        :param period:
        :param symbol:
        :param trade_type:
        :return:
        """
        conn = await self.get_conn()
        df = await conn.fetch("""select open_time, trade_type
         from %s where open_time >= %s and open_time >= %s and open_time < %s and symbol = '%s' and trade_type=%s""" % (
            self.tables[period], self.start * 1000 - 1000, onboard * 1000, self.end * 1000 + 1000, symbol, trade_type))
        # 检测当天是否有数据
        if daily:
            df = set([(x[0] - x[0] % (3600 * 24 * 1000)) for x in df])
            onboard = 0 if len(df) == 0 else min(df) / 1000
            times = set(x * 1000 for x in range(self.start // 86400 * 86400 + 86400, self.end, 86400) if x >= onboard)
        # 按K线周期检测每一条
        else:
            df = set([round(x[0] / period / 1000) * 1000 * period for x in df])
            onboard = 0 if len(df) == 0 else min(df) / 1000
            times = set(
                x * 1000 for x in range(self.start // period * period + period, self.end, period) if x >= onboard)
        comp: Set[int] = times - df
        comp_dt: Set[str] = set()
        if comp:
            comp_dt = set(
                datetime.datetime.fromtimestamp(x / 1000, tz=datetime.timezone.utc).strftime('%Y-%m-%d')
                for x in list(comp))
            logger.error(
                f'表：{self.tables[period]}，币种：\'{symbol}\'{"现货" if trade_type == 0 else "合约"}，'
                f'存在数据：{len(df)}，数据缺失：{comp_dt}')
        else:
            logger.info(
                f'表：{self.tables[period]}币种：\'{symbol}\'，{"现货" if trade_type == 0 else "合约"}存在数据{len(df)}')
        if comp:
            await self.lock.acquire()
            self.info[period][str(trade_type)][symbol] = list(comp_dt) if daily else list(comp)
            self.lock.release()
        await conn.close()

    async def check(self, period: int):
        if self.symbols is None:
            conn = await self.get_conn()
            table = self.tables[period]
            spot_symbols = await conn.fetch("""select distinct symbol from %s 
            where open_time >= %s and open_time < %s and trade_type=0""" % (table, self.start * 1000, self.end * 1000))
            spot_symbols: List[str] = [x[0] for x in spot_symbols]
            swap_symbols = await conn.fetch("""select distinct symbol from %s 
                        where open_time >= %s and open_time < %s and trade_type=1""" % (
                table, self.start * 1000, self.end * 1000))
            swap_symbols: List[str] = [x[0] for x in swap_symbols]
        else:
            spot_symbols = self.symbols
            swap_symbols = self.symbols
        all_task = ([self.check_symbol(period, x, 0) for x in spot_symbols] +
                    [self.check_symbol(period, x, 1) for x in swap_symbols])
        if len(all_task) == 0:
            logger.error(f'表{self.tables[period]}, 时间区间内无数据')
        await asyncio.gather(*all_task)

    def run(self):
        loop = asyncio.new_event_loop()
        for x in self.tables.keys():
            loop.run_until_complete(self.check(x))

        if self.path:
            s_str = datetime.datetime.fromtimestamp(self.start, tz=datetime.timezone.utc).strftime("%Y-%m-%d")
            e_str = datetime.datetime.fromtimestamp(self.end, tz=datetime.timezone.utc).strftime("%Y-%m-%d")
            with open(os.path.join(self.path, f'{s_str}-{e_str}.json'), 'w+') as f:
                json.dump(self.info, f)


if __name__ == '__main__':
    s = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc).timestamp()
    e = datetime.datetime(2024, 8, 3, tzinfo=datetime.timezone.utc).timestamp()
    ch = Checker(int(s), int(e), symbols=['BTCUSDT', 'ETHUSDT'], path=r'/locals/output')
    ch.run()
    # print(ch.info)
